# -*- codeing = utf-8 -*-
# @Time : 2022/2/18 13:30
# @File : data_set.py
# @Software : PyCharm
import config
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder

# 对图像进行预处理

def get_transform():
    # Compose 将多个转化函数组合起来使用
    return transforms.Compose([
        # 图像缩放到32 * 32，opencv抓取的图片尺寸是不一样的
            transforms.Resize(32),
        # 中心裁剪
            transforms.CenterCrop(32),
        # 把像素灰度范围从0-255 变换到0-1
            transforms.ToTensor(),
        # 对每个像素点进行归一化
            transforms.Normalize(mean=[0.4, 0.4, 0.4],
                                 std=[0.2, 0.2, 0.2])
        ])

# 加载训练集
def get_train_dataset(batch_size=10, num_workers=1):
    data_transform = get_transform()
    # ImageFolder 是pytorch框架提供的图片目录加载类  参数：
    # root：图片目录
    # transform: 对PIL Image进行的转换操作，transform的输入是使用loader读取图片的返回对象
    # loader: 给定路径后如何读取图片，默认为RGB格式的PIL Image对象
    train_dataset = ImageFolder(root=config.DATA_TRAIN, transform=data_transform)

    # DataLoader: 生成Pytorch训练时候使用的可迭代对象对象  参数：
    # batch_size : 每次迭代的样本数量
    # shuffle: 打乱顺序
    # loaddata时候使用的进程数量
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    return train_loader# , test_loader

# 加载测试集
def get_test_dataset(batch_size=10, num_workers=1):
    data_transform = get_transform()

    test_dataset = ImageFolder(root=config.DATA_TEST, transform=data_transform)

    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    return test_loader
